package org.fujay.commons.langchain4j.rag.prompt.template;

import dev.langchain4j.rag.content.Content;

import java.util.List;

/**
 * 思维链提示模板
 * 应用思维链(Chain of Thought)方法构建提示，引导模型一步步思考
 */
public class ChainOfThoughtTemplate implements RagPromptTemplate {

    private static final String DEFAULT_TEMPLATE = "你是一个专业的分析助手。请基于检索到的信息回答问题。\n\n" +
            "思考步骤：\n" +
            "1. 分析问题的关键要点\n" +
            "2. 查找相关信息并评估其相关性\n" +
            "3. 推理解答过程\n" +
            "4. 给出完整回答\n" +
            "5. 引用信息来源\n\n" +
            "检索到的信息：\n" +
            "%s\n\n" +
            "用户问题: %s\n\n" +
            "请按以下格式组织你的回答：\n" +
            "分析: [分析问题要点]\n" +
            "思考: [查找信息并推理]\n" +
            "回答: [最终答案]\n" +
            "引用: [列出使用的文档来源，如\"文档[1]\"]\n";

    private final String template;

    /**
     * 默认构造函数
     */
    public ChainOfThoughtTemplate() {
        this(DEFAULT_TEMPLATE);
    }

    /**
     * 构造函数
     *
     * @param template 模板字符串
     */
    public ChainOfThoughtTemplate(String template) {
        this.template = template != null ? template : DEFAULT_TEMPLATE;
    }

    @Override
    public String build(String query, List<Content> retrievedContents) {
        StringBuilder formattedContents = new StringBuilder();

        // 格式化检索内容
        for (int i = 0; i < retrievedContents.size(); i++) {
            Content content = retrievedContents.get(i);
            formattedContents.append("文档[").append(i + 1).append("]: ");
            formattedContents.append(content.toString()).append("\n");

            // 添加来源信息
            var metadata = content.metadata();
            if (metadata != null) {
                String source = getSourceFromMetadata(metadata);
                if (source != null && !source.isEmpty()) {
                    formattedContents.append("来源: ").append(source).append("\n");
                }
            }

            formattedContents.append("\n");
        }

        // 应用模板
        return String.format(template, formattedContents.toString(), query);
    }

    /**
     * 从元数据中获取来源信息
     *
     * @param metadata 元数据
     * @return 来源信息
     */
    private String getSourceFromMetadata(Object metadata) {
        if (metadata == null) {
            return null;
        }

        // 这里简单处理，实际实现可能需要根据元数据类型进行适配
        if (metadata instanceof java.util.Map) {
            java.util.Map<?, ?> map = (java.util.Map<?, ?>) metadata;

            // 尝试获取常见的来源字段
            for (String field : new String[]{"source", "filename", "url", "title", "path"}) {
                Object value = map.get(field);
                if (value != null) {
                    return value.toString();
                }
            }
        }

        return metadata.toString();
    }
}